import torch

from pathlib import Path
# 1. 如果注册了Python算子，则可以通过此方法引入
# from . import _C

#! 注意： 
#! 如果注册了Python算子，则1/2的方式均可；如果只只测了C+/CUDA算子， 则必须使用2；

# 2. 如果只注册了C++/CUDA算子，则必须通过以下方式引入
so_files  = list(Path(__file__).parent.absolute().glob("_C*.so"))
assert len(so_files) == 1, f"Expected 1 so file, but got {len(so_files)}"
torch.ops.load_library(so_files[0])

def muladd(a:torch.Tensor, b:torch.Tensor, c:float):
    return torch.ops._C.muladd(a, b, c)

#! 注册对torch.compile的支持
@torch.library.register_fake("_C::muladd")
def _muladd(a:torch.Tensor, b:torch.Tensor, c:float):
    torch._check(a.dtype  == b.dtype)
    torch._check(a.shape  == b.shape)
    torch._check(a.device == b.device)
    return torch.empty_like(a)


def mul(a:torch.Tensor, b:torch.Tensor):
    return torch.ops._C.mul(a, b)


#! 注册对torch.compile的支持
@torch.library.register_fake("_C::mul")
def _mul(a:torch.Tensor, b:torch.Tensor):
    torch._check(a.dtype  == b.dtype)
    torch._check(a.shape  == b.shape)
    torch._check(a.device == b.device)
    return torch.empty_like(a)



#! c:torch.Tensor是inout的张量
#! 不能将inout(c)张量通过return返回，否则可能与PyTorch的子系统不兼容
def add(a:torch.Tensor, b:torch.Tensor, c:torch.Tensor):
    torch.ops._C.add(a, b, c)



#! 注册对torch.compile的支持
#! 不能将inout(c)张量通过return返回，否则可能与PyTorch的子系统不兼容
@torch.library.register_fake("_C::add")
def _add(a:torch.Tensor, b:torch.Tensor, c:torch.Tensor):
    torch._check(a.dtype  == b.dtype)
    torch._check(a.shape  == b.shape)
    torch._check(a.device == b.device)
